﻿/*
 * @Author: Morphlng
 * @Date: 2022-07-15 11:08:18
 * @LastEditTime: 2022-07-21 17:05:42
 * @LastEditors: Morphlng
 * @Description: A console tool to flip every bit of a file
 * @Compiler-Win: g++ -o3 -std=c++17 main.cpp -o Flipbit
 * @Compiler-Linux: g++ -o3 -std=c++17 -lpthread main.cpp -o Flipbit
 * @FilePath: \Flipbit\main.cpp
 */

#include "argparse.h"
#include "tqdm.h"
#include <fstream>
#include <tuple>
#include <atomic>
#include <thread>

static bool echo{true};

struct MyArgs : public argparse::Args
{
    std::string &positional_src_path = arg("path to the input file").set_default("");
    std::string &positional_out_path = arg("path to the output file").set_default("");
    std::optional<std::string> &explicit_src_path = kwarg("i,input", "path to the input file");
    std::optional<std::string> &explicit_out_path = kwarg("o,output", "path to the output file");
    int &thread_cnt = kwarg("j,thread", "enable multithread if needed").set_default(1);
    bool &split_file = flag("s,split", "A flag to toggle split file reading").set_default(false);
    bool &quiet = flag("q,quiet", "A flag to toggle progress bar").set_default(false);

    void welcome() override
    {
        std::cout << "Welcome to Flipbit!\n";
    }
};

std::tuple<std::string, std::string> parseArgs(const MyArgs &args)
{
    std::string src_path, out_path;

    if (args.explicit_src_path)
    {
        src_path.swap(args.explicit_src_path.value());
    }
    else if (!args.positional_src_path.empty())
    {
        src_path.swap(args.positional_src_path);
    }

    if (args.explicit_out_path)
    {
        out_path.swap(args.explicit_out_path.value());
    }
    else if (!args.positional_out_path.empty())
    {
        out_path.swap(args.positional_out_path);
    }

    if (src_path.empty())
    {
        std::cout << "You must give a valid input file path, exiting!\n";
        exit(1);
    }

    if (out_path.empty())
    {
        out_path = src_path;
        size_t offset = out_path.rfind('.');
        if (offset != std::string::npos)
        {
            out_path.replace(out_path.begin() + offset, out_path.end(), ".rev");
        }
        else
        {
            out_path.append(".rev");
        }
        std::cout << "Using default output path: " << out_path << '\n';
    }

    if (args.quiet)
        echo = false;

    return {src_path, out_path};
}

void flipByte(char &byte)
{
    for (int i = 0; i < 8; i++)
    {
        byte ^= 1 << i;
    }
}

void handleBuffer(std::string &buffer, int threadCnt = 1)
{
    size_t length = buffer.size();

    if (threadCnt > 1)
    {
        std::vector<std::thread> thread_pool(threadCnt);
        if (!echo)
        {
            // no use of atomic, expected performance enhancement
            for (size_t offset = 0; offset < threadCnt; offset++)
            {
                // offset is captured by value
                // otherwise it will be changed while the program is running
                thread_pool[offset] = std::thread([&, offset]()
                                                  {
                    for (size_t i = 0 + offset; i < length; i += threadCnt)
                    {
                        flipByte(buffer[i]);
                    } });
            }
        }
        else
        {
            std::atomic<size_t> tqdm_cnt{0};
            for (size_t offset = 0; offset < threadCnt; offset++)
            {
                thread_pool[offset] = std::thread([&, offset]()
                                                  {
                    for (size_t i = 0 + offset; i < length; i += threadCnt)
                    {
                        flipByte(buffer[i]);
                        tqdm_cnt++;
                    } });
            }

            thread_pool.emplace_back([&tqdm_cnt, length]()
                                     {
                tqdm bar;
                while (tqdm_cnt < length) {
                    bar.progress(tqdm_cnt, length);
                } });
        }
        for (auto &&task : thread_pool)
        {
            task.join();
        }
    }
    else
    {
        if (!echo)
        {
            for (size_t i = 0; i < length; i++)
            {
                flipByte(buffer[i]);
            }
        }
        else
        {
            tqdm bar;
            for (size_t i = 0; i < length; i++)
            {
                bar.progress(i, length);

                flipByte(buffer[i]);
            }
        }
    }
}

int main(int argc, char *argv[])
{
    MyArgs args = argparse::parse<MyArgs>(argc, argv);
    auto [src_path, out_path] = parseArgs(args);

    // Input File Stream
    std::ifstream inFile(src_path, std::ios::binary | std::ios::ate);
    if (!inFile.good())
    {
        std::cout << "Invalid input file, exiting!\n";
        return 1;
    }

    // Output File Stream
    std::ofstream outFile(out_path, std::ios::out | std::ios::binary);

    // Fetch file size
    size_t file_size = inFile.tellg();
    std::cout << "Original file size: " << file_size << " bytes\n";
    inFile.seekg(0);

    // Enable split file reading when file size is bigger than 256MB
    if (!args.split_file && file_size > 268435456)
    {
        args.split_file = true;
        std::cout << "File size is bigger than 256MB, enabling split file read mode\n";
    }

    if (!args.split_file)
    {
        std::string buffer(std::istreambuf_iterator<char>{inFile}, std::istreambuf_iterator<char>{});

        handleBuffer(buffer, args.thread_cnt);

        std::cout << "\nStart to write to file: " << out_path << '\n';
        outFile << buffer;
    }
    else
    {
        std::cout << "\nSplit reading file, 4KB at a time\n";
        // create a 4KB buffer
        std::string buffer(4 * 1024 * 1024, '\0');
        while (inFile)
        {
            inFile.read(buffer.data(), buffer.size());
            size_t readCnt = inFile.gcount();

            /*
                Drop extra bytes when remaining file content
                is smaller than buffer size
            */
            if (readCnt < buffer.size())
            {
                buffer.resize(readCnt);
            }

            handleBuffer(buffer, args.thread_cnt);
            outFile << buffer;
        }
    }

    std::cout << "\nSuccessfully reverse file, wrote to path: " << out_path << '\n';

    return 0;
}